import sys

sys.modules["sqlite3"] = __import__("pysqlite3")  # 劫持旧模块
import chromadb

print('chromadb version:', chromadb.__version__)  # 0.5.x 以上体验最佳

from chromadb import Settings
from chromadb.utils import embedding_functions

# 1. 连库（持久化）
client = chromadb.PersistentClient(
    path="./chroma_db",
    settings=Settings(allow_reset=True))  # allow_reset 方便调试时清库


def prepare_data(client):
    # 2. 选择向量化模型（三选一）
    emb_fn = embedding_functions.OllamaEmbeddingFunction(model_name="nomic-embed-text")

    # 3. 创建/获取集合（类似 MySQL 的表）
    col = client.get_or_create_collection(
        name="faq",
        embedding_function=emb_fn,
        metadata={"hnsw:space": "cosine"})  # 可选：指定距离类型

    # 4. 写入文档（可批量）
    docs = ["Chroma 是一款轻量级向量数据库",
            "RAG 通过检索增强生成，提高大模型回答准确率"]
    metas = [{"src": "intro"}, {"src": "rag"}]
    ids = ["d1", "d2"]

    col.add(documents=docs, metadatas=metas, ids=ids)
    print("已写入", col.count(), "条向量")


def query_data(client):
    col = client.get_or_create_collection(
        name="faq",
    )

    # 5. 语义检索
    ans = col.query(
        query_texts="什么是向量数据库",
        n_results=2,
        where={"src": "intro"})
    print("最相似文本：", ans["documents"])
    print("距离分数：", ans["distances"])  # 越小越相似


query_data(client)
